#include "mlir/IR/BuiltinDialect.h"
#include "myDialect/toyDialect.h"
#include "myDialect/toyOps.h"
#include "rewrite/passes.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/Sequence.h"

using namespace mlir;

struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
using OpRewritePattern<toy::ConstantOp>::OpRewritePattern;
LogicalResult matchAndRewrite(toy::ConstantOp op, PatternRewriter& rewriter) const final {
    auto constValue = op.getInput();
    Location loc = op.getLoc();
    mlir::Value arithop = rewriter.create<arith::ConstantOp>(loc, constValue);
    rewriter.replaceOp(op, arithop);
    return success();
}
};

struct AddOpLowering : public OpRewritePattern<toy::AddOp> {
using OpRewritePattern<toy::AddOp>::OpRewritePattern;
LogicalResult matchAndRewrite(toy::AddOp op, PatternRewriter& rewriter) const final {
    auto left = op.getLhs();
    auto right = op.getRhs();
    Location loc = op.getLoc();
    mlir::Value arithOp = rewriter.create<arith::AddIOp>(loc, left, right);
    rewriter.replaceOp(op, arithOp);
    return success();
}
};

namespace {
struct ToyToArithLoweringPass : public PassWrapper<ToyToArithLoweringPass, OperationPass<ModuleOp>> {
    MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToArithLoweringPass)

    void getDependentDialects(DialectRegistry &registry) const override {
        registry.insert<arith::ArithDialect>();
    }

    void runOnOperation() final;
};
}

void ToyToArithLoweringPass::runOnOperation() {
    ConversionTarget target(this->getContext());
    target.addLegalDialect<arith::ArithDialect, BuiltinDialect>();
    target.addIllegalDialect<toy::ToyDialect>();
    target.addDynamicallyLegalOp<toy::PrintOp>(
        [](toy::PrintOp op){
            return op.getOperand().getType().isa<IntegerType>();
        });
    RewritePatternSet patterns(&getContext());
    patterns.add<ConstantOpLowering, AddOpLowering>(&getContext());

    if (failed(
        applyPartialConversion(getOperation(), target, std::move(patterns))))
    signalPassFailure();
}

std::unique_ptr<Pass> mlir::toy::createLowerToArithPass() {
    return std::make_unique<ToyToArithLoweringPass>();
}






